import os
import pandas as pd
import folder_paths
import torch
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import numpy as np
from io import BytesIO
import os

class Read_EXCEL:
    """读取Excel文件并匹配节目数据
    读取节目表和数据表，根据节目名匹配数据并返回图片列表
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "programs_file_path": ("STRING", {"default": "", "multiline": False, "placeholder": "例如: C:/Users/name/Documents/programs.xlsx 或 programs.xlsx"}),
                "data_file_path": ("STRING", {"default": "", "multiline": False, "placeholder": "例如: C:/Users/name/Documents/data.xlsx 或 data.xlsx"})
            },
            "optional": {
                "programs_sheet_name": ("STRING", {"default": "Sheet1", "visible": False}),
                "data_sheet_name": ("STRING", {"default": "Sheet1", "visible": False}),
                "program_name_column": ("STRING", {"default": "节目名", "visible": False}),
                "template_requirement_column": ("STRING", {"default": "模版要求", "visible": False}),
                "ip_address_column": ("STRING", {"default": "ip文件地址", "visible": False}),
                "background_address_column": ("STRING", {"default": "背景地址", "visible": False}),
                "logo_address_column": ("STRING", {"default": "logo图片", "visible": False}),
                "image_naming_pattern": ("STRING", {"default": "{program_name}_{width}x{height}", "placeholder": "例如: {program_name}_{width}x{height}", "tooltip": "使用{program_name}、{width}、{height}等占位符"})
            }
        }

    RETURN_TYPES = ("LIST", "LIST", "LIST", "LIST", "LIST", "INT")
    RETURN_NAMES = ("节目名列表", "ip图片列表", "背景图片列表", "logo图片列表", "图片名称列表", "数量")
    FUNCTION = "process_excel_data"
    CATEGORY = "psd"

    def process_excel_data(self, programs_file_path, data_file_path, programs_sheet_name, data_sheet_name, program_name_column, template_requirement_column, ip_address_column, background_address_column, logo_address_column, image_naming_pattern="{program_name}_{width}x{height}"):
        # 规范化文件路径（处理不同操作系统的路径分隔符）
        programs_file_path = os.path.normpath(programs_file_path)
        data_file_path = os.path.normpath(data_file_path)

        # 获取完整文件路径
        try:
            if os.path.isabs(programs_file_path):
                programs_full_path = programs_file_path
            else:
                # 尝试相对路径（相对于ComfyUI的输入目录）
                programs_full_path = os.path.join(folder_paths.get_input_directory(), programs_file_path)
                # 如果在输入目录中找不到，尝试相对于当前工作目录
                if not os.path.exists(programs_full_path):
                    programs_full_path = os.path.abspath(programs_file_path)

            if os.path.isabs(data_file_path):
                data_full_path = data_file_path
            else:
                # 尝试相对路径（相对于ComfyUI的输入目录）
                data_full_path = os.path.join(folder_paths.get_input_directory(), data_file_path)
                # 如果在输入目录中找不到，尝试相对于当前工作目录
                if not os.path.exists(data_full_path):
                    data_full_path = os.path.abspath(data_file_path)

            # 检查文件是否存在
            if not os.path.exists(programs_full_path):
                error_msg = f"错误: 节目表文件不存在于路径: {programs_full_path}"
                # 确保空图像的数据类型正确
                empty_image = torch.zeros((1, 3, 256, 256), dtype=torch.float32)
                return ([error_msg], [empty_image], [empty_image], [empty_image], [], 0)
            if not os.path.exists(data_full_path):
                error_msg = f"错误: 数据表文件不存在于路径: {data_full_path}"
                # 确保空图像的数据类型正确
                empty_image = torch.zeros((1, 3, 256, 256), dtype=torch.float32)
                return ([error_msg], [empty_image], [empty_image], [empty_image], [], 0)

            # 检查文件扩展名是否为Excel文件
            if not (programs_full_path.lower().endswith('.xlsx') or programs_full_path.lower().endswith('.xls')):
                error_msg = f"错误: 节目表文件不是有效的Excel文件: {programs_full_path}"
                # 确保空图像的数据类型正确
                empty_image = torch.zeros((1, 3, 256, 256), dtype=torch.float32)
                return ([error_msg], [empty_image], [empty_image], [empty_image], [], 0)
            if not (data_full_path.lower().endswith('.xlsx') or data_full_path.lower().endswith('.xls')):
                error_msg = f"错误: 数据表文件不是有效的Excel文件: {data_full_path}"
                # 确保空图像的数据类型正确
                empty_image = torch.zeros((1, 3, 256, 256), dtype=torch.float32)
                return ([error_msg], [empty_image], [empty_image], [empty_image], [], 0)
        except Exception as e:
            error_msg = f"路径处理错误: {str(e)}"
            # 确保空图像的数据类型正确
            empty_image = torch.zeros((1, 3, 256, 256), dtype=torch.float32)
            return ([error_msg], [empty_image], [empty_image], [empty_image], [], 0)

        try:
            # 读取节目表文件
            programs_df = pd.read_excel(programs_full_path, sheet_name=programs_sheet_name)
            # 读取数据表文件
            data_df = pd.read_excel(data_full_path, sheet_name=data_sheet_name)

            # 数据预处理 - 去除空格和标准化字符串
            for df in [programs_df, data_df]:
                if program_name_column in df.columns:
                    df[program_name_column] = df[program_name_column].astype(str).str.strip().str.lower()
                if template_requirement_column in df.columns:
                    df[template_requirement_column] = df[template_requirement_column].astype(str).str.strip().str.lower()

            # 检查必要的列是否存在
            required_columns = [
                (programs_df, program_name_column, "节目表"),
                (programs_df, template_requirement_column, "节目表"),
                (data_df, program_name_column, "数据表"),
                (data_df, ip_address_column, "数据表"),
                (data_df, background_address_column, "数据表"),
                (data_df, logo_address_column, "数据表")
            ]

            for df, col, table_name in required_columns:
                if col not in df.columns:
                    error_msg = f"错误: {table_name}中找不到'{col}'列"
                    # 确保空图像的数据类型正确
                    empty_image = torch.zeros((1, 3, 256, 256), dtype=torch.float32)
                    return ([error_msg], [empty_image], [empty_image], [empty_image], [], 0)

            # 处理每一行数据
            results = []
            for _, program_row in programs_df.iterrows():
                program_name = program_row.get(program_name_column, "")
                template_requirement = program_row.get(template_requirement_column, "")

                if not program_name:
                    continue

                # 精确匹配节目名
                print(f"正在匹配节目: {program_name}")
                matched_data_rows = data_df[data_df[program_name_column] == program_name]
                if matched_data_rows.empty:
                    print(f"未找到匹配的节目数据: {program_name}")
                    # 确保空图像的数据类型正确
                    empty_image = torch.zeros((1, 3, 256, 256), dtype=torch.float32)
                    results.append((str(program_name), empty_image, empty_image, empty_image))
                    continue
                else:
                    print(f"找到 {len(matched_data_rows)} 条匹配数据")

                # 获取匹配行的数据
                data_row = matched_data_rows.iloc[0]
                ip_address = data_row.get(ip_address_column, "")
                background_address = data_row.get(background_address_column, "")
                logo_address = data_row.get(logo_address_column, "")

                # 标准化模板要求
                template_requirement = str(template_requirement).lower()

                # 根据模版要求筛选需要输出的地址
                print(f"模板要求: {template_requirement}")
                ip_output = str(ip_address) if ("修改角色ip" in template_requirement or "ip" in template_requirement) else None
                background_output = str(background_address) if "背景" in template_requirement else None
                logo_output = str(logo_address) if "logo" in template_requirement else None
                print(f"筛选结果 - IP: {ip_output}, 背景: {background_output}, Logo: {logo_output}")

                # 如果没有指定特定要求，输出所有地址
                if not template_requirement or (ip_output is None and background_output is None and logo_output is None):
                    ip_output = str(ip_address)
                    background_output = str(background_address)
                    logo_output = str(logo_address)

                # 确保非空字符串
                ip_output = ip_output or ""
                background_output = background_output or ""
                logo_output = logo_output or ""

                # 加载图片
                ip_image = self.load_image(ip_output)
                background_image = self.load_image(background_output)
                logo_image = self.load_image(logo_output)
                
                # 验证所有图像的格式
                print(f"IP图片 - 形状: {ip_image.shape}, 数据类型: {ip_image.dtype}")
                print(f"背景图片 - 形状: {background_image.shape}, 数据类型: {background_image.dtype}")
                print(f"Logo图片 - 形状: {logo_image.shape}, 数据类型: {logo_image.dtype}")
                
                # 最终验证，确保所有图像都是正确的格式
                assert ip_image.shape[1] in [3, 4], f"IP图片不是3或4通道: {ip_image.shape}"
                assert background_image.shape[1] in [3, 4], f"背景图片不是3或4通道: {background_image.shape}"
                assert logo_image.shape[1] in [3, 4], f"Logo图片不是3或4通道: {logo_image.shape}"
                assert ip_image.dtype == torch.float32, f"IP图片数据类型错误: {ip_image.dtype}"
                assert background_image.dtype == torch.float32, f"背景图片数据类型错误: {background_image.dtype}"
                assert logo_image.dtype == torch.float32, f"Logo图片数据类型错误: {logo_image.dtype}"
                
                results.append((str(program_name), ip_image, background_image, logo_image))

            cleaned_results = results
            print(f"数据处理完成，共生成 {len(cleaned_results)} 条结果")

            # 最终验证所有返回的图像
            for result in cleaned_results:
                program_name, ip_img, bg_img, logo_img = result
                
                # 验证IP图片
                if ip_img.shape[1] not in [3, 4] or ip_img.dtype != torch.float32:
                    print(f"警告: {program_name}的IP图片格式不正确! 形状: {ip_img.shape}, 数据类型: {ip_img.dtype}")
                    # 强制修正
                    if ip_img.shape[1] not in [3, 4]:
                        ip_img = ip_img.repeat(1, 3, 1, 1) if ip_img.shape[1] == 1 else ip_img[:, :3, :, :]
                    if ip_img.dtype != torch.float32:
                        ip_img = ip_img.to(torch.float32)
                    print(f"已强制修正为: 形状: {ip_img.shape}, 数据类型: {ip_img.dtype}")
                
                # 验证背景图片
                if bg_img.shape[1] not in [3, 4] or bg_img.dtype != torch.float32:
                    print(f"警告: {program_name}的背景图片格式不正确! 形状: {bg_img.shape}, 数据类型: {bg_img.dtype}")
                    # 强制修正
                    if bg_img.shape[1] not in [3, 4]:
                        bg_img = bg_img.repeat(1, 3, 1, 1) if bg_img.shape[1] == 1 else bg_img[:, :3, :, :]
                    if bg_img.dtype != torch.float32:
                        bg_img = bg_img.to(torch.float32)
                    print(f"已强制修正为: 形状: {bg_img.shape}, 数据类型: {bg_img.dtype}")
                
                # 验证Logo图片
                if logo_img.shape[1] not in [3, 4] or logo_img.dtype != torch.float32:
                    print(f"警告: {program_name}的Logo图片格式不正确! 形状: {logo_img.shape}, 数据类型: {logo_img.dtype}")
                    # 强制修正
                    if logo_img.shape[1] not in [3, 4]:
                        logo_img = logo_img.repeat(1, 3, 1, 1) if logo_img.shape[1] == 1 else logo_img[:, :3, :, :]
                    if logo_img.dtype != torch.float32:
                        logo_img = logo_img.to(torch.float32)
                    print(f"已强制修正为: 形状: {logo_img.shape}, 数据类型: {logo_img.dtype}")
                
                # 更新结果列表中的当前项
                idx = cleaned_results.index(result)
                cleaned_results[idx] = (program_name, ip_img, bg_img, logo_img)
            
            # 最终安全检查：确保所有返回的图像都符合ComfyUI要求
            def final_sanitize_image(img):
                print(f"进入最终安全检查 - 初始形状: {img.shape}, 数据类型: {img.dtype}")
                # 允许3或4通道
                if img.shape[1] not in [3, 4]:
                    print(f"最终安全检查: 修正通道数从 {img.shape[1]} 到 3")
                    img = img.repeat(1, 3, 1, 1) if img.shape[1] == 1 else img[:, :3, :, :]
                    print(f"通道修正后 - 形状: {img.shape}")
                # 确保float32数据类型
                if img.dtype != torch.float32:
                    print(f"最终安全检查: 修正数据类型从 {img.dtype} 到 float32")
                    img = img.to(torch.float32)
                    print(f"数据类型修正后 - 数据类型: {img.dtype}")
                # 添加额外的转换步骤，确保数据格式完全符合ComfyUI要求
                img_np = img.cpu().numpy()  # 转换为NumPy数组
                print(f"转换为NumPy后 - 形状: {img_np.shape}, 数据类型: {img_np.dtype}")
                # 允许3或4通道
                if img_np.shape[1] not in [3, 4]:
                    print(f"NumPy数组通道数错误，修正从 {img_np.shape[1]} 到 3")
                    if img_np.shape[1] == 1:
                        img_np = np.repeat(img_np, 3, axis=1)
                    else:
                        img_np = img_np[:, :3, :, :]
                    print(f"NumPy通道修正后 - 形状: {img_np.shape}")
                # 确保是float32
                img_np = img_np.astype(np.float32)  # 确保是float32
                print(f"NumPy类型转换后 - 数据类型: {img_np.dtype}")
                # 转换回PyTorch张量
                img = torch.from_numpy(img_np).to(img.device)  # 转换回PyTorch张量
                print(f"转换回PyTorch后 - 形状: {img.shape}, 数据类型: {img.dtype}")
                # 确保连续性
                if not img.is_contiguous():
                    img = img.contiguous()
                    print(f"确保连续性后 - 形状: {img.shape}, 数据类型: {img.dtype}")
                # 最终验证
                # 允许3或4通道图像
                # 允许3或4通道图像
                assert img.shape[1] in [3, 4], f"最终验证失败: 图像不是3或4通道，形状为 {img.shape}"
                assert img.dtype == torch.float32, f"最终验证失败: 图像数据类型不是float32，而是 {img.dtype}"
                print(f"最终安全检查完成 - 形状: {img.shape}, 数据类型: {img.dtype}")
                return img

            # 处理所有结果，应用最终安全检查
            program_names = []
            ip_images = []
            bg_images = []
            logo_images = []
            image_names = []
            
            for result in cleaned_results:
                program_name, ip_img, bg_img, logo_img = result
                
                # 应用最终安全检查
                ip_img = final_sanitize_image(ip_img)
                bg_img = final_sanitize_image(bg_img)
                logo_img = final_sanitize_image(logo_img)
                
                print(f"最终返回 - {program_name} - IP图片: {ip_img.shape}, {ip_img.dtype}")
                print(f"最终返回 - {program_name} - 背景图片: {bg_img.shape}, {bg_img.dtype}")
                print(f"最终返回 - {program_name} - Logo图片: {logo_img.shape}, {logo_img.dtype}")
                
                # 生成图片名称
                # 使用IP图片的尺寸作为参考
                width = ip_img.shape[3]
                height = ip_img.shape[2]
                # 替换占位符
                image_name = image_naming_pattern.format(
                    program_name=program_name,
                    width=width,
                    height=height
                )
                # 确保文件名有效
                image_name = image_name.replace('/', '_').replace('\\', '_').replace(':', '_')
                
                program_names.append(program_name)
                ip_images.append(ip_img)
                bg_images.append(bg_img)
                logo_images.append(logo_img)
                image_names.append(image_name)
            
            # 返回列表形式的结果
            return (program_names, ip_images, bg_images, logo_images, image_names, len(cleaned_results))

        except Exception as e:
            error_msg = f"处理错误: {str(e)}"
            # 返回空图片，确保数据类型正确
            empty_image = torch.zeros((1, 3, 256, 256), dtype=torch.float32)
            return ([error_msg], [empty_image], [empty_image], [empty_image], [], 0)

    def is_valid_image_format(self, image_path):
        """验证图片格式是否支持"""
        valid_extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff']
        ext = os.path.splitext(image_path.lower())[1]
        return ext in valid_extensions

    def load_image(self, image_path): 
        """增强版的图片加载函数，解决各种异常情况"""
        from io import BytesIO 
        import os
        
        # 1. 检查路径是否存在且是文件 
        if not image_path or not os.path.exists(image_path) or not os.path.isfile(image_path): 
            print(f"无效的图片路径: {image_path}") 
            placeholder = self.get_placeholder_image("路径无效或不是文件")
            print(f"返回占位图片，形状: {placeholder.shape}, 数据类型: {placeholder.dtype}")
            return placeholder
        
        # 2. 验证图片格式
        if not self.is_valid_image_format(image_path):
            print(f"不支持的图片格式: {image_path}")
            placeholder = self.get_placeholder_image("不支持的图片格式")
            print(f"返回占位图片，形状: {placeholder.shape}, 数据类型: {placeholder.dtype}")
            return placeholder
            
        # 3. 尝试读取文件内容 
        try: 
            with open(image_path, 'rb') as f: 
                file_content = f.read() 
                
            # 4. 验证是否为有效图片（至少有一定大小） 
            if len(file_content) < 128:  # 小于128字节不是有效图片 
                placeholder = self.get_placeholder_image("文件过小")
                print(f"返回占位图片，形状: {placeholder.shape}, 数据类型: {placeholder.dtype}")
                return placeholder
                
            # 5. 使用BytesIO安全解码 
            img = Image.open(BytesIO(file_content)) 
            
            # 6. 保留原始通道信息（包括alpha通道）
            if img.mode == 'RGBA':
                # 保留RGBA格式
                pass
            elif img.mode == 'LA':
                # 转换为RGBA
                img = img.convert('RGBA')
            else:
                # 其他模式转换为RGB
                img = img.convert('RGB') 
            
            # 7. 添加更健壮的形状验证 
            if img.width < 1 or img.height < 1: 
                raise ValueError("图片尺寸无效") 
            
            # 8. 调整大小（保持原比例） 
            min_size = 64 
            if img.width < min_size or img.height < min_size: 
                scale = min_size / min(img.width, img.height) 
                new_width = max(int(img.width * scale), min_size) 
                new_height = max(int(img.height * scale), min_size) 
                img = img.resize((new_width, new_height), Image.LANCZOS) 
            
            # 9. 转换为Tensor - 确保数据类型正确
            img_array = np.array(img, dtype=np.uint8)  # 先转为uint8
            # 添加调试信息：检查数组值范围
            print(f"图像数组值范围 (uint8): [{np.min(img_array)}, {np.max(img_array)}]")
            img_array = img_array.astype(np.float32) / 255.0  # 再归一化到[0,1]
            # 添加调试信息：检查归一化后的值范围
            print(f"图像数组值范围 (归一化后): [{np.min(img_array)}, {np.max(img_array)}]")
            img_tensor = torch.from_numpy(img_array).permute(2, 0, 1)  # 转为CHW格式
            img_tensor = img_tensor.unsqueeze(0).contiguous()  # 添加批次维度并确保连续性
            
            # 额外验证：确保数据类型正确，允许3或4通道
            if img_tensor.dtype != torch.float32: 
                print(f"警告: 图片格式不正确! 形状: {img_tensor.shape}, 数据类型: {img_tensor.dtype}") 
                # 强制修正数据类型
                img_tensor = img_tensor.to(torch.float32) 
                if img_tensor.dtype != torch.float32: 
                    # 转换为float32
                    img_tensor = img_tensor.to(torch.float32) 
                print(f"已强制修正数据类型为: {img_tensor.dtype}") 
            
            # 验证输出格式和值范围
            print(f"成功加载图片: {image_path}, 形状: {img_tensor.shape}, 数据类型: {img_tensor.dtype}")
            print(f"图像张量值范围: [{torch.min(img_tensor)}, {torch.max(img_tensor)}]")
            return img_tensor 
            
        except Exception as e: 
            print(f"无法加载图片 '{image_path}': {str(e)}") 
            placeholder = self.get_placeholder_image(str(e))
            print(f"返回占位图片，形状: {placeholder.shape}, 数据类型: {placeholder.dtype}")
            return placeholder 

    def get_placeholder_image(self, message): 
        """创建带错误信息的占位图片"""
        # 创建纯灰色背景 (带alpha通道)
        placeholder = Image.new('RGBA', (256, 256), (128, 128, 128, 255)) 
        
        # 添加错误信息 
        try: 
            from PIL import ImageDraw, ImageFont 
            draw = ImageDraw.Draw(placeholder) 
            
            # 尝试加载字体 
            try: 
                font = ImageFont.truetype("arial.ttf", 20) 
            except: 
                font = ImageFont.load_default() 
                
            draw.text((10, 110), message, fill=(255, 0, 0, 255), font=font) 
        except: 
            pass  # 忽略文本渲染错误 
        
        # 转换为Tensor - 确保数据类型正确
        img_array = np.array(placeholder, dtype=np.uint8)  # 先转为uint8
        img_array = img_array.astype(np.float32) / 255.0  # 再归一化到[0,1]
        img_tensor = torch.from_numpy(img_array).permute(2, 0, 1)  # 转为CHW格式
        img_tensor = img_tensor.unsqueeze(0).contiguous()  # 添加批次维度并确保连续性
        
        # 额外验证：确保数据类型正确，允许3或4通道
        if img_tensor.dtype != torch.float32: 
            print(f"警告: 占位图片数据类型不正确! 数据类型: {img_tensor.dtype}") 
            # 强制修正数据类型
            img_tensor = img_tensor.to(torch.float32) 
            print(f"已强制修正数据类型为: {img_tensor.dtype}") 
        
        # 确保通道数为3或4
        if img_tensor.shape[1] not in [3, 4]:
            print(f"警告: 占位图片通道数不正确! 通道数: {img_tensor.shape[1]}")
            # 转换为3通道
            if img_tensor.shape[1] == 1:
                img_tensor = img_tensor.repeat(1, 3, 1, 1)
            else:
                img_tensor = img_tensor[:, :3, :, :]
            print(f"已强制修正通道数为: {img_tensor.shape[1]}")
        
        return img_tensor 

    def is_valid_image_format(self, img_path): 
        """验证是否为支持的有效图片格式""" 
        valid_extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff'] 
        return any(img_path.lower().endswith(ext) for ext in valid_extensions)

# 注册节点
NODE_CLASS_MAPPINGS = {
    "Read_EXCEL": Read_EXCEL
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Read_EXCEL": "📊 读取Excel数据"
}